# -*- coding: utf-8 -*-
#
# Modified by Peize Sun, Rufeng Zhang
# Contact: {sunpeize, cxrfzhang}@foxmail.com
#
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
# from detectron2.config import CfgNode as CN
from yacs.config import CfgNode as CN

_C = CN()
_C.CUDNN_BENCHMARK = False
_C.VERSION = 2
_C._BASE_ = ''
_C.TRAIN_PRINT_FREQ = 20

_C.DATALOADER = CN()
_C.DATALOADER.ASPECT_RATIO_GROUPING = True
_C.DATALOADER.FILTER_EMPTY_ANNOTATIONS = False
_C.DATALOADER.NUM_WORKERS = 4
_C.DATALOADER.REPEAT_THRESHOLD = 0.0
_C.DATALOADER.SAMPLER_TRAIN = 'TrainingSampler'

_C.DATASETS = CN()
_C.DATASETS.PRECOMPUTED_PROPOSAL_TOPK_TEST = 1000
_C.DATASETS.PRECOMPUTED_PROPOSAL_TOPK_TRAIN = 2000
_C.DATASETS.PROPOSAL_FILES_TEST = ()
_C.DATASETS.PROPOSAL_FILES_TRAIN = ()
_C.DATASETS.TEST = ('coco_2017_val',)
_C.DATASETS.TRAIN = ('coco_2017_train',)

_C.GLOBAL = CN()
_C.GLOBAL.HACK = 1.0

_C.INPUT = CN()
_C.INPUT.CROP = CN()
_C.INPUT.CROP.ENABLED = False
_C.INPUT.CROP.SIZE = [384, 600]
_C.INPUT.CROP.TYPE = 'absolute_range'

_C.INPUT.FORMAT = 'RGB'
_C.INPUT.MASK_FORMAT = 'polygon'
_C.INPUT.MAX_SIZE_TEST = 1333
_C.INPUT.MAX_SIZE_TRAIN = 1333
_C.INPUT.MIN_SIZE_TEST = 800 #768
_C.INPUT.MIN_SIZE_TRAIN = (480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800)
_C.INPUT.MIN_SIZE_TRAIN_SAMPLING = 'choice'
_C.INPUT.RANDOM_FLIP = 'horizontal'

_C.MODEL = CN()
_C.MODEL.BACKBONE = CN()
_C.MODEL.BACKBONE.FREEZE_AT = 2
_C.MODEL.BACKBONE.NAME = 'build_resnet_fpn_backbone'
_C.MODEL.DEVICE = 'cuda'
_C.MODEL.META_ARCHITECTURE = 'SparseRCNN'
_C.MODEL.PIXEL_MEAN = [123.675, 116.28, 103.53]
_C.MODEL.PIXEL_STD = [58.395, 57.12, 57.375]

_C.MODEL.RESNETS = CN()
_C.MODEL.RESNETS.DEFORM_MODULATED = False
_C.MODEL.RESNETS.DEFORM_NUM_GROUPS = 1
_C.MODEL.RESNETS.DEFORM_ON_PER_STAGE = [False, False, False, False]
_C.MODEL.RESNETS.DEPTH = 50
_C.MODEL.RESNETS.NORM = 'FrozenBN'
_C.MODEL.RESNETS.NUM_GROUPS = 1
_C.MODEL.RESNETS.OUT_FEATURES = ['res2', 'res3', 'res4', 'res5']
_C.MODEL.RESNETS.RES2_OUT_CHANNELS = 256
_C.MODEL.RESNETS.RES5_DILATION = 1
_C.MODEL.RESNETS.STEM_OUT_CHANNELS = 64
_C.MODEL.RESNETS.STRIDE_IN_1X1 = False
_C.MODEL.RESNETS.WIDTH_PER_GROUP = 64

_C.MODEL.FPN = CN()
_C.MODEL.FPN.FUSE_TYPE = 'sum'
_C.MODEL.FPN.IN_FEATURES = ['res2', 'res3', 'res4', 'res5']
_C.MODEL.FPN.NORM = ''
_C.MODEL.FPN.OUT_CHANNELS = 256

_C.MODEL.ROI_BOX_HEAD = CN()
_C.MODEL.ROI_BOX_HEAD.USE_LEVEL_POOLER = True
_C.MODEL.ROI_BOX_HEAD.BBOX_REG_LOSS_TYPE = 'smooth_l1'
_C.MODEL.ROI_BOX_HEAD.BBOX_REG_LOSS_WEIGHT = 1.0
_C.MODEL.ROI_BOX_HEAD.BBOX_REG_WEIGHTS = (10.0, 10.0, 5.0, 5.0)
_C.MODEL.ROI_BOX_HEAD.CLS_AGNOSTIC_BBOX_REG = False
_C.MODEL.ROI_BOX_HEAD.CONV_DIM = 256
_C.MODEL.ROI_BOX_HEAD.FC_DIM = 1024
_C.MODEL.ROI_BOX_HEAD.NUM_CONV = 0
_C.MODEL.ROI_BOX_HEAD.NUM_FC = 0
_C.MODEL.ROI_BOX_HEAD.POOLER_RESOLUTION = 7
_C.MODEL.ROI_BOX_HEAD.POOLER_SAMPLING_RATIO = 2
_C.MODEL.ROI_BOX_HEAD.POOLER_TYPE = 'ROIAlignV2'
_C.MODEL.ROI_BOX_HEAD.SMOOTH_L1_BETA = 0.0
_C.MODEL.ROI_BOX_HEAD.TRAIN_ON_PRED_BOXES = False

_C.MODEL.ROI_HEADS = CN()
_C.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 512
_C.MODEL.ROI_HEADS.IN_FEATURES = ['p2', 'p3', 'p4', 'p5']
_C.MODEL.ROI_HEADS.IOU_LABELS = [0, 1]
_C.MODEL.ROI_HEADS.IOU_THRESHOLDS = [0.5]
_C.MODEL.ROI_HEADS.NAME = 'Res5ROIHeads'
_C.MODEL.ROI_HEADS.NMS_THRESH_TEST = 0.5
_C.MODEL.ROI_HEADS.NUM_CLASSES = 80
_C.MODEL.ROI_HEADS.POSITIVE_FRACTION = 0.25
_C.MODEL.ROI_HEADS.PROPOSAL_APPEND_GT = True
_C.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.05

_C.MODEL.SparseRCNN = CN()
_C.MODEL.SparseRCNN.ACTIVATION = 'relu'
_C.MODEL.SparseRCNN.ALPHA = 0.25
_C.MODEL.SparseRCNN.CLASS_WEIGHT = 2.0
_C.MODEL.SparseRCNN.DEEP_SUPERVISION = True
_C.MODEL.SparseRCNN.DIM_DYNAMIC = 64
_C.MODEL.SparseRCNN.DIM_FEEDFORWARD = 2048
_C.MODEL.SparseRCNN.DROPOUT = 0.0
_C.MODEL.SparseRCNN.GAMMA = 2.0
_C.MODEL.SparseRCNN.GIOU_WEIGHT = 2.0
_C.MODEL.SparseRCNN.HIDDEN_DIM = 256
_C.MODEL.SparseRCNN.L1_WEIGHT = 5.0
_C.MODEL.SparseRCNN.NHEADS = 8
_C.MODEL.SparseRCNN.NO_OBJECT_WEIGHT = 0.1
_C.MODEL.SparseRCNN.NUM_CLASSES = 80
_C.MODEL.SparseRCNN.NUM_CLS = 1
_C.MODEL.SparseRCNN.NUM_DYNAMIC = 2
_C.MODEL.SparseRCNN.NUM_HEADS = 6
_C.MODEL.SparseRCNN.NUM_PROPOSALS = 100
_C.MODEL.SparseRCNN.NUM_REG = 3
_C.MODEL.SparseRCNN.PRIOR_PROB = 0.01
_C.MODEL.SparseRCNN.USE_FOCAL = True

_C.MODEL.WEIGHTS = 'models/R-50.pkl'
_C.OUTPUT_DIR = './output'
_C.SEED = 40244023
_C.BASE_ROOT = './datasets/coco'

_C.SOLVER = CN()
_C.SOLVER.AMP = CN()
_C.SOLVER.AMP.ENABLED = False
_C.SOLVER.BACKBONE_MULTIPLIER = 1.0
_C.SOLVER.BASE_LR = 2.5e-05
_C.SOLVER.BIAS_LR_FACTOR = 1.0
_C.SOLVER.CHECKPOINT_PERIOD = 5000
_C.SOLVER.CLIP_GRADIENTS = CN()
_C.SOLVER.CLIP_GRADIENTS.CLIP_TYPE = 'full_model'
_C.SOLVER.CLIP_GRADIENTS.CLIP_VALUE = 1.0
_C.SOLVER.CLIP_GRADIENTS.ENABLED = True
_C.SOLVER.CLIP_GRADIENTS.NORM_TYPE = 2.0
_C.SOLVER.GAMMA = 0.1
_C.SOLVER.IMS_PER_BATCH = 5 #6
_C.SOLVER.LR_SCHEDULER_NAME = 'WarmupMultiStepLR'
_C.SOLVER.MAX_ITER = 270000
_C.SOLVER.MOMENTUM = 0.9
_C.SOLVER.NESTEROV = False
_C.SOLVER.OPTIMIZER = 'ADAMW'
_C.SOLVER.REFERENCE_WORLD_SIZE = 0
_C.SOLVER.STEPS = (210000, 250000)
_C.SOLVER.WARMUP_FACTOR = 0.01
_C.SOLVER.WARMUP_ITERS = 1000
_C.SOLVER.WARMUP_METHOD = 'linear'
_C.SOLVER.WEIGHT_DECAY = 0.0001
_C.SOLVER.WEIGHT_DECAY_BIAS = 0.0001
_C.SOLVER.WEIGHT_DECAY_NORM = 0.0

_C.TEST = CN()
_C.TEST.AUG = CN()
_C.TEST.AUG.ENABLED = False
_C.TEST.AUG.FLIP = True
_C.TEST.AUG.MAX_SIZE = 4000
_C.TEST.AUG.MIN_SIZES = (400, 500, 600, 700, 800, 900, 1000, 1100, 1200)
_C.TEST.DETECTIONS_PER_IMAGE = 100
_C.TEST.EVAL_PERIOD = 7330
_C.TEST.EXPECTED_RESULTS = []
_C.TEST.KEYPOINT_OKS_SIGMAS = []
_C.TEST.PRECISE_BN = CN()
_C.TEST.PRECISE_BN.ENABLED = False
_C.TEST.PRECISE_BN.NUM_ITER = 200


def add_sparsercnn_config(cfg):
    """
    Add config for SparseRCNN.
    """
    cfg.MODEL.SparseRCNN = CN()
    cfg.MODEL.SparseRCNN.NUM_CLASSES = 80
    cfg.MODEL.SparseRCNN.NUM_PROPOSALS = 300

    # RCNN Head.
    cfg.MODEL.SparseRCNN.NHEADS = 8
    cfg.MODEL.SparseRCNN.DROPOUT = 0.0
    cfg.MODEL.SparseRCNN.DIM_FEEDFORWARD = 2048
    cfg.MODEL.SparseRCNN.ACTIVATION = 'relu'
    cfg.MODEL.SparseRCNN.HIDDEN_DIM = 256
    cfg.MODEL.SparseRCNN.NUM_CLS = 1
    cfg.MODEL.SparseRCNN.NUM_REG = 3
    cfg.MODEL.SparseRCNN.NUM_HEADS = 6

    # Dynamic Conv.
    cfg.MODEL.SparseRCNN.NUM_DYNAMIC = 2
    cfg.MODEL.SparseRCNN.DIM_DYNAMIC = 64

    # Loss.
    cfg.MODEL.SparseRCNN.CLASS_WEIGHT = 2.0
    cfg.MODEL.SparseRCNN.GIOU_WEIGHT = 2.0
    cfg.MODEL.SparseRCNN.L1_WEIGHT = 5.0
    cfg.MODEL.SparseRCNN.DEEP_SUPERVISION = True
    cfg.MODEL.SparseRCNN.NO_OBJECT_WEIGHT = 0.1

    # Focal Loss.
    cfg.MODEL.SparseRCNN.USE_FOCAL = True
    cfg.MODEL.SparseRCNN.ALPHA = 0.25
    cfg.MODEL.SparseRCNN.GAMMA = 2.0
    cfg.MODEL.SparseRCNN.PRIOR_PROB = 0.01

    # Optimizer.
    cfg.SOLVER.OPTIMIZER = "ADAMW"
    cfg.SOLVER.BACKBONE_MULTIPLIER = 1.0
